/*
 * This is written by Lulu Han.
 * E-mail: locomotive_crypto@163.com
 */

#include <iostream>
#include <exception>
#include <boost/program_options.hpp>
#include "elgamal.h"

namespace op = boost::program_options;

void test_elgamal_for_correctness(long mod_bitlens, long group_bitlens);

int main(int argc, char* argv[]){
	op::options_description desc("All options for testing Elgamal algorithm");
	desc.add_options()
		("help", "Produce help message")
		("modbits", op::value<long>(), "Modulus length")
		("groupbits", op::value<long>(), "Group size for discrete logarithm problem");
	
	op::variables_map vm;
	op::store(op::parse_command_line(argc, argv, desc), vm);
	op::notify(vm);
	
	if(argc == 1 || vm.count("help")){
		std::cout << desc << std::endl;
		return 0;
	}
		
	if(vm.count("modbits") && vm.count("groupbits")){
		test_elgamal_for_correctness(
			vm["modbits"].as<long>(), 
				vm["groupbits"].as<long>());
		return 0;
	}
	return 0;
}

void test_elgamal_for_correctness(long mod_bitlens, long group_bitlens){
	Elgamal::SecretKey sk;
	Elgamal::PublicKey pk;
	Elgamal::key_gen(sk, pk, mod_bitlens, group_bitlens);
	
	Elgamal::Encryptor enc(pk);
	Elgamal::Decryptor dec(sk);
	
	NTL::ZZ m;
	m = NTL::RandomLen_ZZ(32);
	
	HEnc::PTxt pt;
	HEnc::CTxt ct;
	
	pt.set_pt(m);
	std::cout << "pt(m) = " << pt.get_pt() << std::endl;
	
	// Encrypt a message pt
	enc.encrypt(pt, ct);
	std::cout << "ct(m): "<< std::endl;
	std::cout << "\t c1 = " << ct.get_ct1() << std::endl;
	std::cout << "\t c2 = " << ct.get_ct2() << std::endl;
	
	// Decrypt a message ct
	HEnc::PTxt rt;
	dec.decrypt(ct, rt);
	std::cout << "rt(m) = " << rt.get_pt() << std::endl;
	
	if(rt.get_pt() != pt.get_pt()){
		throw std::runtime_error("Testing encrypt() and decrypt() functions failed...");
	}
	
	// Testing for homomorphic multiplication
	// 1) Homomorphic multiplication for two ciphertexts
	std::cout << "\nHomomorphic multiplication for two ciphertexts..." << std::endl;
	NTL::ZZ m2;
	m2 = NTL::RandomLen_ZZ(32);
	
	HEnc::PTxt pt2;
	pt2.set_pt(m2);
	std::cout << "pt2 = " << pt2.get_pt() << std::endl;
	
	HEnc::CTxt ct2;
	enc.encrypt(pt2, ct2);
	
	enc.he_mul(ct, ct2, ct);
	std::cout << "ct(pt * pt2): "<< std::endl;
	std::cout << "\t c1 = " << ct.get_ct1() << std::endl;
	std::cout << "\t c2 = " << ct.get_ct2() << std::endl;
	
	dec.decrypt(ct, rt);
	std::cout << "rt(ct * ct2) = " << rt.get_pt() << std::endl;
	
	if(rt.get_pt() != (pt.get_pt() * pt2.get_pt())){
		throw std::runtime_error("Testing he_mul() functions failed...");
	}
	
	// 2) Homomorphic multiplication for one ciphertext and one plaintext
	std::cout << "\nHomomorphic multiplication for one ciphertext and one plaintext..." << std::endl;
	NTL::ZZ m3;
	m3 = NTL::RandomLen_ZZ(32);
	
	HEnc::PTxt pt3;
	pt3.set_pt(m3);
	std::cout << "pt3 = " << pt3.get_pt() << std::endl;
	
	enc.he_mul(ct2, pt3, ct);
	std::cout << "ct(ct2 * pt3): "<< std::endl;
	std::cout << "\t c1 = " << ct.get_ct1() << std::endl;
	std::cout << "\t c2 = " << ct.get_ct2() << std::endl;
	
	dec.decrypt(ct, rt);
	std::cout << "rt(pt2 * pt3) = " << rt.get_pt() << std::endl;
	
	if(rt.get_pt() != (pt2.get_pt() * pt3.get_pt())){
		throw std::runtime_error("Testing he_mul() functions failed...");
	}
}
